#include "hnurm_armor/armor_node.hpp"
#include <cv_bridge/cv_bridge.h>

#include <sensor_msgs/image_encodings.hpp>
#include <tf2/LinearMath/Quaternion.h>
#include <tf2_geometry_msgs/tf2_geometry_msgs.hpp>

using namespace std::chrono_literals;

namespace hnurm
{
    ArmorNode::ArmorNode(const rclcpp::NodeOptions &options)
        : Node("armor_node", options)
    {
        std::string tracker_info_topic = this->declare_parameter("tracker_info_topic", "tracker_info");
        std::string target_topic = this->declare_parameter("target_topic", "target");
        std::string vision_send_topic = this->declare_parameter("vision_send_topic", "vision_send_data");

        RCLCPP_INFO(get_logger(), "ArmorNode is running");
        processor_ = std::make_unique<Processor>(shared_from_this());
        RCLCPP_INFO(get_logger(), "Processor created");
        compensator_ = std::make_unique<Compensator>(shared_from_this());
        RCLCPP_INFO(get_logger(), "Compensator created");

        send_data_pub_
            = this->create_publisher<hnurm_interfaces::msg::VisionSendData>(vision_send_topic, rclcpp::ServicesQoS());

        tracker_info_pub_
            = this->create_publisher<hnurm_interfaces::msg::TrackerInfo>(tracker_info_topic, rclcpp::SensorDataQoS());
        target_pub_ = this->create_publisher<hnurm_interfaces::msg::Target>(target_topic, rclcpp::SensorDataQoS());

        armor_array_sub_ = this->create_subscription<hnurm_interfaces::msg::ArmorArray>(
            "armor", rclcpp::SensorDataQoS(), std::bind(&ArmorNode::armor_array_callback, this, std::placeholders::_1)
        );

        auto using_timer = this->declare_parameter("using_timer", false);
        if (using_timer)
            timer_ = this->create_wall_timer(1ms, std::bind(&ArmorNode::timer_callback, this));

        init_markers();
    }
    float ArmorNode::calculateAverage(const std::vector<float> &vec)
    {
        int n = vec.size();
        if (n == 0) return 0.0f;

        float sum = 0.0f;
        for (float value: vec)
        {
            sum += value;
        }
        return sum / n;
    }
    void ArmorNode::Averagedistance(std::vector<float> &p, float value)
    {
        if (p.size() >= 50)
        {
            p.erase(p.begin());
        }
        p.push_back(value);
    }
    void ArmorNode::timer_callback()
    {
        std::lock_guard<std::mutex> lock(mtx);
        TargetInfo tmp_target;

        // update processor and tracker
        processor_->ProcessArmor(armors, last_recv_array_.recv_uart, tmp_target);

        // send target info over uart
        send_target(tmp_target, last_recv_array_.recv_uart);
    }

    void ArmorNode::init_markers()
    {
        position_marker_.ns = "position";
        position_marker_.type = visualization_msgs::msg::Marker::SPHERE;
        position_marker_.scale.x = position_marker_.scale.y = position_marker_.scale.z = 0.1;
        position_marker_.color.a = 1.0;
        position_marker_.color.g = 1.0;
        linear_v_marker_.type = visualization_msgs::msg::Marker::ARROW;
        linear_v_marker_.ns = "linear_v";
        linear_v_marker_.scale.x = 0.03;
        linear_v_marker_.scale.y = 0.05;
        linear_v_marker_.color.a = 1.0;
        linear_v_marker_.color.r = 1.0;
        linear_v_marker_.color.g = 1.0;
        angular_v_marker_.type = visualization_msgs::msg::Marker::ARROW;
        angular_v_marker_.ns = "angular_v";
        angular_v_marker_.scale.x = 0.03;
        angular_v_marker_.scale.y = 0.05;
        angular_v_marker_.color.a = 1.0;
        angular_v_marker_.color.b = 1.0;
        angular_v_marker_.color.g = 1.0;
        armor_marker_.ns = "armors";
        armor_marker_.type = visualization_msgs::msg::Marker::CUBE;
        armor_marker_.scale.x = 0.03;
        armor_marker_.scale.z = 0.125;
        armor_marker_.color.a = 1.0;
        armor_marker_.color.r = 1.0;
        std::string topic = this->declare_parameter("marker_topic", "/tracker/marker");
        marker_pub_ = this->create_publisher<visualization_msgs::msg::MarkerArray>(topic, 10);
    }

    void ArmorNode::armor_array_callback(const hnurm_interfaces::msg::ArmorArray::SharedPtr msg)
    {
        std::lock_guard<std::mutex> lock(mtx);
        last_recv_array_ = *msg;

        TargetInfo tmp_target;

        armors = fromROSMsg(*msg);

        // update processor and tracker
        processor_->ProcessArmor(armors, msg->recv_uart, tmp_target);

        // send target info over uart
        send_target(tmp_target, msg->recv_uart);

        // publish info and markers
        // tracker info
        hnurm_interfaces::msg::TrackerInfo info_msg;
        hnurm_interfaces::msg::Target target_msg;
        target_msg.header.stamp = msg->header.stamp;
        target_msg.header.frame_id = "odom";

        info_msg.position_diff = processor_->tracker->info_position_diff;
        info_msg.yaw_diff = processor_->tracker->info_yaw_diff;
        info_msg.position.x = processor_->tracker->measurement(0);
        info_msg.position.y = processor_->tracker->measurement(1);
        info_msg.position.z = processor_->tracker->measurement(2);
        info_msg.yaw = processor_->tracker->measurement(3);
        info_msg.yaw = angles::to_degrees(info_msg.yaw);
        info_msg.distance = processor_->tracker->measurement.norm();
        tracker_info_pub_->publish(info_msg);

        if (processor_->tracker->tracker_state == Tracker::DETECTING)
        {
            target_msg.tracking = false;
        }
        else if (processor_->tracker->tracker_state == Tracker::TRACKING
            || processor_->tracker->tracker_state == Tracker::TEMP_LOST)
        {
            target_msg.tracking = true;
            // Fill target message
            const auto &state = processor_->tracker->target_state;
            Averagedistance(pv, state(7));
            target_msg.id = processor_->tracker->tracked_id;
            target_msg.armors_num = static_cast<int>(processor_->tracker->tracked_armors_num);
            target_msg.position.x = state(0);
            target_msg.velocity.x = state(1);
            target_msg.position.y = state(2);
            target_msg.velocity.y = state(3);
            target_msg.position.z = state(4);
            target_msg.velocity.z = state(5);
            target_msg.yaw = state(6);
            target_msg.v_yaw = calculateAverage(pv);
            target_msg.radius_1 = state(8);
            target_msg.radius_2 = processor_->tracker->another_r;
            target_msg.dz = processor_->tracker->dz;
            target_msg.armor_pitch = processor_->tracker->armor_pitch;
            target_msg.armor_yaw = processor_->tracker->armor_yaw;
        }
        if (!armors.empty())
        {
            target_pub_->publish(target_msg);
        }

        publish_markers(target_msg);
    }

    void ArmorNode::send_target(TargetInfo &target, hnurm_interfaces::msg::VisionRecvData &uart_msg)
    {
        hnurm_interfaces::msg::VisionSendData send;
        send.header.frame_id = "serial_send";
        send.header.stamp = this->now();
        send.control_id = uart_msg.control_id;  // copy the control id
        if (processor_->tracker->tracker_state != Tracker::LOST && processor_->tracker->tracker_state != Tracker::JUMP)
        {
            float armor_yaw =
                compensator_->CalcFinalAngle(target, uart_msg, send, processor_->tracker->tracked_armors_num);
            if (abs(send.yaw - uart_msg.yaw) < 3 && abs(send.pitch - uart_msg.pitch) < 1.3 && armor_yaw < 45)
            {
                send.target_state.data = hnurm_interfaces::msg::TargetState::FIRE;
            }
            else
            {
                send.target_state.data = hnurm_interfaces::msg::TargetState::CONVERGING;
            }
            send.target_distance = Eigen::Vector3f(target.x, target.y, target.z).norm();
            send_data_pub_->publish(send);
        }
        else
        {
            send.target_state.data = hnurm_interfaces::msg::TargetState::LOST_TARGET;
            send_data_pub_->publish(send);
        }
    }

    void ArmorNode::publish_markers(hnurm_interfaces::msg::Target &target_msg)
    {
        position_marker_.header = target_msg.header;
        linear_v_marker_.header = target_msg.header;
        angular_v_marker_.header = target_msg.header;
        armor_marker_.header = target_msg.header;

        visualization_msgs::msg::MarkerArray marker_array;
        if (target_msg.tracking)
        {
            double yaw = target_msg.yaw, r1 = target_msg.radius_1, r2 = target_msg.radius_2;
            double xc = target_msg.position.x, yc = target_msg.position.y, za = target_msg.position.z;
            double vx = target_msg.velocity.x, vy = target_msg.velocity.y, vz = target_msg.velocity.z;
            double dz = target_msg.dz;

            position_marker_.action = visualization_msgs::msg::Marker::ADD;
            position_marker_.pose.position.x = xc;
            position_marker_.pose.position.y = yc;
            position_marker_.pose.position.z = za + dz / 2;

            linear_v_marker_.action = visualization_msgs::msg::Marker::ADD;
            linear_v_marker_.points.clear();
            linear_v_marker_.points.emplace_back(position_marker_.pose.position);
            geometry_msgs::msg::Point arrow_end = position_marker_.pose.position;
            arrow_end.x += vx;
            arrow_end.y += vy;
            arrow_end.z += vz;
            linear_v_marker_.points.emplace_back(arrow_end);

            angular_v_marker_.action = visualization_msgs::msg::Marker::ADD;
            angular_v_marker_.points.clear();
            angular_v_marker_.points.emplace_back(position_marker_.pose.position);
            arrow_end = position_marker_.pose.position;
            arrow_end.z += target_msg.v_yaw / M_PI;
            angular_v_marker_.points.emplace_back(arrow_end);

            armor_marker_.action = visualization_msgs::msg::Marker::ADD;
            armor_marker_.scale.y = processor_->tracker->tracked_armor.armor_type == ArmorType::SMALL ? 0.135 : 0.23;
            bool is_current_pair = true;
            size_t a_n = target_msg.armors_num;
            geometry_msgs::msg::Point p_a;
            double r = 0;
            for (size_t i = 0; i < a_n; i++)
            {
                double tmp_yaw = yaw + i * (2 * M_PI / a_n);
                // Only 4 armors has 2 radius and height
                if (a_n == 4)
                {
                    r = is_current_pair ? r1 : r2;
                    p_a.z = za + (is_current_pair ? 0 : dz);
                    is_current_pair = !is_current_pair;
                }
                else
                {
                    r = r1;
                    p_a.z = za;
                }
                p_a.x = xc - r * cos(tmp_yaw);
                p_a.y = yc - r * sin(tmp_yaw);

                armor_marker_.id = i;
                armor_marker_.pose.position = p_a;
                tf2::Quaternion q;
                q.setRPY(0, target_msg.id == "outpost" ? -0.26 : 0.26, tmp_yaw);
                armor_marker_.pose.orientation = tf2::toMsg(q);
                marker_array.markers.emplace_back(armor_marker_);
            }
        }
        else
        {
            position_marker_.action = visualization_msgs::msg::Marker::DELETE;
            linear_v_marker_.action = visualization_msgs::msg::Marker::DELETE;
            angular_v_marker_.action = visualization_msgs::msg::Marker::DELETE;

            armor_marker_.action = visualization_msgs::msg::Marker::DELETE;
            marker_array.markers.emplace_back(armor_marker_);
        }

        marker_array.markers.emplace_back(position_marker_);
        marker_array.markers.emplace_back(linear_v_marker_);
        marker_array.markers.emplace_back(angular_v_marker_);
        marker_pub_->publish(marker_array);
    }

    std::vector<Armor> fromROSMsg(const hnurm_interfaces::msg::ArmorArray &msg)
    {
        std::vector<Armor> armors;
        armors.resize(msg.armors.size());
        for (size_t i = 0; i < msg.armors.size(); i++)
        {
            armors[i].left_light.color = msg.armors[i].left_light.color;
            armors[i].left_light.top.x = msg.armors[i].left_light.top.x;
            armors[i].left_light.top.y = msg.armors[i].left_light.top.y;
            armors[i].left_light.center.x = msg.armors[i].left_light.center.x;
            armors[i].left_light.center.y = msg.armors[i].left_light.center.y;
            armors[i].left_light.bottom.x = msg.armors[i].left_light.bottom.x;
            armors[i].left_light.bottom.y = msg.armors[i].left_light.bottom.y;
            armors[i].left_light.size.width = msg.armors[i].left_light.size.x;
            armors[i].left_light.size.height = msg.armors[i].left_light.size.y;
            armors[i].left_light.length = msg.armors[i].left_light.length;
            armors[i].left_light.width = msg.armors[i].left_light.width;
            armors[i].left_light.tilt_angle = msg.armors[i].left_light.tilt_angle;
            armors[i].left_light.angle = msg.armors[i].left_light.angle;

            armors[i].right_light.color = msg.armors[i].right_light.color;
            armors[i].right_light.top.x = msg.armors[i].right_light.top.x;
            armors[i].right_light.top.y = msg.armors[i].right_light.top.y;
            armors[i].right_light.center.x = msg.armors[i].right_light.center.x;
            armors[i].right_light.center.y = msg.armors[i].right_light.center.y;
            armors[i].right_light.bottom.x = msg.armors[i].right_light.bottom.x;
            armors[i].right_light.bottom.y = msg.armors[i].right_light.bottom.y;
            armors[i].right_light.size.width = msg.armors[i].right_light.size.x;
            armors[i].right_light.size.height = msg.armors[i].right_light.size.y;
            armors[i].right_light.length = msg.armors[i].right_light.length;
            armors[i].right_light.width = msg.armors[i].right_light.width;
            armors[i].right_light.tilt_angle = msg.armors[i].right_light.tilt_angle;
            armors[i].right_light.angle = msg.armors[i].right_light.angle;

            armors[i].center.x = msg.armors[i].center.x;
            armors[i].center.y = msg.armors[i].center.y;

            armors[i].points2d.resize(msg.armors[i].points2d.size());
            for (size_t j = 0; j < msg.armors[i].points2d.size(); j++)
            {
                armors[i].points2d[j].x = msg.armors[i].points2d[j].x;
                armors[i].points2d[j].y = msg.armors[i].points2d[j].y;
            }

            armors[i].number_img
                = cv_bridge::toCvCopy(msg.armors[i].number_img, sensor_msgs::image_encodings::MONO8)->image;
            armors[i].number = msg.armors[i].number;
            armors[i].idx = msg.armors[i].idx;
            armors[i].similarity = msg.armors[i].similarity;
            armors[i].confidence = msg.armors[i].confidence;
            armors[i].classification_result = msg.armors[i].classification_result;
            armors[i].armor_type = ArmorType(msg.armors[i].armor_type.data);
        }

        return armors;
    }

}  // namespace hnurm